import matplotlib.pyplot as plt
import torch
import numpy as np


def show_image(im):
    fix_image_processing()

    plt.figure(figsize=(5,5), dpi=120)
    plt.imshow(im)
    plt.axis('off')
    plt.show()


def fix_image_processing():
    import os
    os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"


# batch = [num_of_masks, x, y]
def merge_masks(batch):
    return torch.sum(batch, dim=0)


# batch = [num_of_masks, x, y]
def merge_masks_with_colors(batch):
    color = 1
    for mask in batch:
        mask[mask > 0] = color
        color += 1
    return torch.sum(batch, dim=0)


def merge_image_and_masks(image, masks):
    merged_masks_with_color_dim = torch \
        .stack((masks,masks,masks)) \
        .permute(1,2,0)

    stacked_example = np.hstack((merged_masks_with_color_dim, image))
    return stacked_example


# boxes - torch.Size([num_of_masks, 4])
def build_box_masks(x, y, boxes):
    box_mask = np.zeros((len(boxes), x, y))
    counter = 0
    for box in boxes:
        xmin, ymin, xmax, ymax = box
        box_mask[counter, ymin.long():ymax.long(), xmin.long():xmax.long()] = counter + 1
        counter += 1
    box_mask = torch.tensor(box_mask)
    return box_mask


# image - torch.Size([x, y, 3])
# masks - torch.Size([x, y])
# box_mask_view - torch.Size([x, y, 3])
def merge_image_and_masks_boxes(image, masks, box_mask_view):

    merged_masks_with_color_dim = torch \
        .stack((masks,masks,masks)) \
        .permute(1,2,0)

    merged_box_masks_with_color_dim = torch \
        .stack((box_mask_view,box_mask_view,box_mask_view)) \
        .permute(1,2,0)

    stacked_example = np.hstack((
        merged_masks_with_color_dim,
        image,
        merged_box_masks_with_color_dim))
    return stacked_example


def show_image_grid(rows, cols, image_batch):
    img_count = 0
    fig, axes = plt.subplots(nrows=rows, ncols=cols, figsize=(15, 15))
    for i in range(rows):
        for j in range(cols):
            if img_count < len(image_batch):
                axes[i, j].imshow(image_batch[img_count].permute(1, 2, 0))
                axes[i, j].axis('off')
                img_count += 1
